import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import datasets
from torchvision.transforms import ToTensor, Compose, Normalize, Resize
from torch.utils.data import DataLoader
from opacus import PrivacyEngine
from tqdm import tqdm

ENABLE_DP = True
BATCH_SIZE = 256
TEST_BATCH_SIZE = 1024
EPOCHS = 10
LR = 0.1
SIGMA = 1.0
C = 1.0
DELTA = 1e-5
device = "cuda" if torch.cuda.is_available() else "cpu"

train_dataset = datasets.MNIST(
    root="./mnist",
    train=True,
    download=True,
    transform=Compose([
        Resize(10),
        ToTensor(),
        # Normalize((0.1307,), (0.3081,))
    ])
)
test_dataset = datasets.MNIST(
    root="./mnist",
    train=False,
    download=True,
    transform=Compose([
        Resize(10),
        ToTensor(),
        # Normalize((0.1307,), (0.3081,))
    ])
)
train_dataloader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset, batch_size=TEST_BATCH_SIZE, shuffle=False)


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(100, 20)
        self.fc2 = nn.Linear(20, 15)
        self.fc3 = nn.Linear(15, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.fc1(x)
        x = F.leaky_relu(x)
        x = self.fc2(x)
        x = F.leaky_relu(x)
        x = self.fc3(x)
        x = F.leaky_relu(x)
        return x


model = Model().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LR)
if ENABLE_DP:
    privacy_engine = PrivacyEngine()
    model, optimizer, train_dataloader = privacy_engine.make_private(
        module=model,
        optimizer=optimizer,
        data_loader=train_dataloader,
        noise_multiplier=SIGMA,
        max_grad_norm=C
    )

vec_corr_per_epoch = [[] for i in range(3)]  # 向量相关性
dim_corr_per_epoch = [[] for i in range(3)]  # 维度相关性


def train(epoch):
    model.train()
    losses = []
    vec_corr_fc1 = np.zeros((20, 20))
    vec_corr_fc2 = np.zeros((15, 15))
    vec_corr_fc3 = np.zeros((10, 10))
    i = 0
    flat_grads_fc1 = []
    flat_grads_fc2 = []
    flat_grads_fc3 = []
    correct = 0
    total = 0
    for batch_id, (data, target) in enumerate(tqdm(train_dataloader)):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        losses.append(loss.item())
        pred = output.argmax(
            dim=1, keepdim=True
        )
        correct += pred.eq(target.view_as(pred)).sum().item()
        total += data.shape[0]
        loss.backward()
        i += 1
        for name, param in model.named_parameters():
            if name in ['fc1.weight', '_module.fc1.weight']:
                grad = param.grad.cpu().numpy()
                vec_corr_fc1 += np.abs(np.corrcoef(grad))
                flat_grads_fc1.append([x for x in grad.reshape(-1)])
            if name in ['fc2.weight', '_module.fc2.weight']:
                grad = param.grad.cpu().numpy()
                vec_corr_fc2 += np.abs(np.corrcoef(grad))
                flat_grads_fc2.append([x for x in grad.reshape(-1)])
            if name in ['fc3.weight', '_module.fc3.weight']:
                grad = param.grad.cpu().numpy()
                vec_corr_fc3 += np.abs(np.corrcoef(grad))
                flat_grads_fc3.append([x for x in grad.reshape(-1)])
        optimizer.step()
    vec_corr_fc1 /= i
    vec_corr_fc2 /= i
    vec_corr_fc3 /= i
    vec_corr_per_epoch[0].append(vec_corr_fc1)
    vec_corr_per_epoch[1].append(vec_corr_fc2)
    vec_corr_per_epoch[2].append(vec_corr_fc3)
    dim_corr_fc1 = np.abs(np.corrcoef(flat_grads_fc1, rowvar=False))
    dim_corr_fc2 = np.abs(np.corrcoef(flat_grads_fc2, rowvar=False))
    dim_corr_fc3 = np.abs(np.corrcoef(flat_grads_fc3, rowvar=False))
    dim_corr_per_epoch[0].append(dim_corr_fc1)
    dim_corr_per_epoch[1].append(dim_corr_fc2)
    dim_corr_per_epoch[2].append(dim_corr_fc3)
    accuracy = correct / total
    print(f"Train Epoch: {epoch + 1} \t Loss: {np.mean(losses):.6f} \t Accuracy: {accuracy:.4f}")


def test():
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for data, target in tqdm(test_dataloader):
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += loss_fn(output, target).item()
            pred = output.argmax(
                dim=1, keepdim=True
            )
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_dataloader.dataset)
    print(
        "Test Set: \t Loss: {:.4f} \t Accuracy: {}/{} ({:.2f}%)".format(
            test_loss,
            correct,
            len(test_dataloader.dataset),
            100.0 * correct / len(test_dataloader.dataset),
        )
    )


for i in range(EPOCHS):
    train(i)
test()
print("Done!")
